/*
 * Copyright(C) 2022. Huawei Technologies Co.,Ltd. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <iostream>
#include <map>
#include <fstream>
#include "unistd.h"
#include <memory>
#include <queue>
#include <thread>
#include "boost/filesystem.hpp"

#include "MxBase/DeviceManager/DeviceManager.h"
#include "MxBase/DvppWrapper/DvppWrapper.h"
#include "MxBase/MemoryHelper/MemoryHelper.h"

#include "MxBase/E2eInfer/ImageProcessor/ImageProcessor.h"
#include "MxBase/E2eInfer/VideoDecoder/VideoDecoder.h"
#include "MxBase/E2eInfer/VideoEncoder/VideoEncoder.h"
#include "MxBase/E2eInfer/DataType.h"

#include "MxBase/MxBase.h"
#include "MxBase/Log/Log.h"

#include "BlockingQueue.h"

extern "C"
{
#include <libavformat/avformat.h>
}

namespace
{
    using namespace MxBase;
    const uint32_t SRC_WIDTH = 1920;
    const uint32_t SRC_HEIGHT = 1080;
    const uint32_t MAX_WIDTH = 3840;
    const uint32_t MAX_HEIGHT = 2160;
    const uint32_t MAX_FRAME_COUNT = 5000;
    const uint32_t TARGET_FRAME_COUNT = 400;
    const uint32_t MS_TIMEOUT = 2000;
    const uint32_t DEVICE_ID = 0;
    const uint32_t CHANNEL_ID = 1;
    const uint32_t TIME_QUEUE = 500; // ms
    const uint32_t EOS_WAITE = 5;
}
namespace fs = boost::filesystem;

struct FrameImage
{
    Image image; // video Image Class
    uint32_t frameId = 0;
    uint32_t channelId = 0;
};
namespace
{
    // 内存队列
    BlockingQueue<FrameImage> frameImageQueue;
    AVFormatContext *pFormatCtx = nullptr;
    BlockingQueue<FrameImage> encodeQueue;
    
    bool stop_flag = false;
}

// ffmpeg 拉流
AVFormatContext *CreateFormatContext(std::string filePath)
{
    LogInfo << "start to CreateFormatContext!";
    // creat message for stream pull
    AVFormatContext *formatContext = nullptr;
    AVDictionary *options = nullptr;

    LogInfo << "start to avformat_open_input!";
    int ret = avformat_open_input(&formatContext, filePath.c_str(), nullptr, &options);
    if (options != nullptr)
    {
        av_dict_free(&options);
    }
    if (ret != 0)
    {
        LogError << "Couldn`t open input stream " << filePath.c_str() << " ret=" << ret;
        return nullptr;
    }
    ret = avformat_find_stream_info(formatContext, nullptr);
    if (ret != 0)
    {
        LogError << "Couldn`t open input stream information";
        return nullptr;
    }
    return formatContext;
}

// 获取H264中的帧
void GetFrame(AVPacket &pkt, FrameImage &frameImage, AVFormatContext *pFormatCtx)
{
    av_init_packet(&pkt);
    int ret = av_read_frame(pFormatCtx, &pkt);
    if (ret != 0)
    {
        LogInfo << "[StreamPuller] channel Read frame failed, continue!";
        if (ret = +AVERROR_EOF)
        {
            LogInfo << "[StreamPuller] channel StreamPuller is EOF, over!";
            return;
        }
        return;
    }
    else
    {
        if (pkt.size <= 0)
        {
            LogError << "Invalid pkt.size: " << pkt.size;
            return;
        }

        // send to the device
        auto hostDeleter = [](void *dataPtr) -> void {};
        MemoryData data(pkt.size, MemoryData::MEMORY_HOST);
        MemoryData src((void *)(pkt.data), pkt.size, MemoryData::MEMORY_HOST);
        APP_ERROR ret = MemoryHelper::MxbsMallocAndCopy(data, src);
        if (ret != APP_ERR_OK)
        {
            LogError << "MxbsMallocAndCopy failed!";
        }
        std::shared_ptr<uint8_t> imageData((uint8_t *)data.ptrData, hostDeleter);

        Image subImage(imageData, pkt.size);
        frameImage.image = subImage;

        LogDebug << "'channelId = " << frameImage.channelId << ", frameId = " << frameImage.frameId << " , dataSize = " << frameImage.image.GetDataSize();

        av_packet_unref(&pkt);
    }
    return;
}

// 视频解码回调 userData为videoDecoder->Decode()传入的自定义数据，参阅文档：API参考（C++） V2>类参考>VideoDecoder>Decode
APP_ERROR CallBackVdec(Image &decodedImage, uint32_t channelId, uint32_t frameId, void *userData)
{
    FrameImage frameImage;
    frameImage.image = decodedImage;
    frameImage.channelId = channelId;
    frameImage.frameId = frameId;

    BlockingQueue<FrameImage> *p_frameImageQueue = (BlockingQueue<FrameImage> *)userData;
    p_frameImageQueue->Push(frameImage);
    LogInfo << "frameId(" << frameImage.frameId << ") decoded successfully.";
    return APP_ERR_OK;
}

// 视频编码回调 userData为videoEncoder->Encode()传入的自定义数据，参阅文档：API参考（C++） V2>类参考>VideoEncoder>Encode
APP_ERROR CallBackVenc(std::shared_ptr<uint8_t> &outDataPtr, uint32_t &outDataSize, uint32_t &channelId,
                       uint32_t &frameId, void *userData)
{
    Image image(outDataPtr, outDataSize, -1, Size(MAX_WIDTH, MAX_HEIGHT));
    FrameImage frameImage;
    frameImage.image = image;
    frameImage.channelId = channelId;
    frameImage.frameId = frameId;

    BlockingQueue<FrameImage> *p_encodeQueue = (BlockingQueue<FrameImage> *)userData;
    p_encodeQueue->Push(frameImage);
    LogInfo << "frameId(" << frameImage.frameId << ") encoded successfully.";
    return APP_ERR_OK;
}

// 视频流解码线程
void VdecThread(size_t frameCount)
{
    AVPacket pkt;
    uint32_t frameId = 0;
    int32_t channelId = CHANNEL_ID;
    // 解码器参数
    VideoDecodeConfig config;
    VideoDecodeCallBack cPtr = CallBackVdec;
    config.width = SRC_WIDTH;
    config.height = SRC_HEIGHT;
    config.callbackFunc = cPtr;
    config.skipInterval = 0; // 跳帧控制

    std::shared_ptr<VideoDecoder> videoDecoder = std::make_shared<VideoDecoder>(config, DEVICE_ID, channelId);
    while (!stop_flag)
    {
        if (frameId >= TARGET_FRAME_COUNT)
        {
            return;
        }
        Image subImage;
        FrameImage frame;
        frame.channelId = 0;
        frame.frameId = frameId;
        frame.image = subImage;

        GetFrame(pkt, frame, pFormatCtx);
        APP_ERROR ret = videoDecoder->Decode(frame.image.GetData(), frame.image.GetDataSize(), frameId, &frameImageQueue);
        if (ret != APP_ERR_OK)
        {
            LogError << "videoDecoder Decode failed. ret is: " << ret;
        }
        frameId++;
    }
}

// 图像编码h264线程
void VencThread()
{
    // 编码器参数
    VideoEncodeConfig vEConfig;
    VideoEncodeCallBack cEPtr = CallBackVenc;
    vEConfig.callbackFunc = cEPtr;
    vEConfig.width = 1280;
    vEConfig.height = 720;
    // 用户可自定义编码参数
    vEConfig.keyFrameInterval = 50;
    vEConfig.srcRate = 25;
    vEConfig.rcMode = 0;
    vEConfig.maxBitRate = 2000;
    vEConfig.ipProp = 50;
    std::shared_ptr<VideoEncoder> videoEncoder = std::make_shared<VideoEncoder>(vEConfig, DEVICE_ID);

    int frameCount = 0;
    while (!stop_flag)
    {
        if (frameCount >= TARGET_FRAME_COUNT)
        {
            return;
        }
        
        if (frameImageQueue.GetSize() > 0)
        {
            FrameImage frameImage;
            APP_ERROR ret = frameImageQueue.Pop(frameImage, TIME_QUEUE);
            if (ret != APP_ERR_OK)
            {
                LogError << "frameImageQueue pop failed" << ret;
                continue;
            }
            videoEncoder->Encode(frameImage.image, frameImage.frameId, &encodeQueue);
            frameCount += 1;
        }
    }
}

void WriteThread(std::string saveFilePath)
{
    FILE *fp = fopen(saveFilePath.c_str(), "wb");
    if (fp == nullptr)
    {
        LogError << "Failed to open file.";
        return;
    }

    int frameCount = 0;
    int empty_count = 0;
    while (!stop_flag)
    {
        if (empty_count > EOS_WAITE || frameCount >= TARGET_FRAME_COUNT)
        {
            stop_flag = true;
            LogInfo << "process and write frame:" << frameCount;
        }
        
        FrameImage encodeTemp;
        APP_ERROR ret = encodeQueue.Pop(encodeTemp, TIME_QUEUE);
        if (ret != APP_ERR_OK)
        {
            LogError << "encodeQueue pop failed:" << ret;
            if (encodeQueue.GetSize() == 0)
            {
                empty_count++;
            }
            continue;
        }

        std::shared_ptr<uint8_t> data_sp = encodeTemp.image.GetData();
        void *data_p = data_sp.get();
        if (fwrite(data_p, encodeTemp.image.GetDataSize(), 1, fp) != 1)
        {
            LogInfo << "write frame to file fail";
        }

        frameCount++;
    }
}

void PullStream(std::string filePath)
{
    // av_register_all();
    avformat_network_init();
    pFormatCtx = avformat_alloc_context();
    pFormatCtx = CreateFormatContext(filePath);
    av_dump_format(pFormatCtx, 0, filePath.c_str(), 0);
}

int main(int argc, char *argv[])
{
    if (argc <= 1)
    {
        LogWarn << "Please input image path, such as './video_sample test.264'.";
        return APP_ERR_OK;
    }
    APP_ERROR ret;

    // global init
    ret = MxInit();
    if (ret != APP_ERR_OK)
    {
        LogError << "MxInit failed, ret=" << ret << ".";
    }

    std::string filePath = argv[1];
    std::string outPath = "out_" + filePath;
    PullStream(filePath);
    std::thread threadVdec(VdecThread, TARGET_FRAME_COUNT);
    std::thread threadVenc(VencThread);
    std::thread threadWrite(WriteThread, outPath);

    threadVdec.join();
    threadVenc.join();
    threadWrite.join();

    LogInfo << "ALL DONE";
}